import os
import argparse

FILE_ENDWITH = "_HybrdChip"

def compare_files(mask_file_path, readout_file_path, output_file_path, special_value):
    # 打开掩码文件和读出文件
    with open(mask_file_path, 'r') as mask_file, open(readout_file_path, 'r') as readout_file:
        # 逐行读取文件内容
        mask_lines = mask_file.readlines()
        readout_lines = readout_file.readlines()
        count = 0

    # 创建一个新文件来保存修改后的读出文件内容
    with open(output_file_path, 'w') as output_file:
        for mask_line, readout_line in zip(mask_lines, readout_lines):
            # 假设每行都是由空格分隔的二进制值
            mask_values = mask_line.strip()
            readout_values = readout_line.strip()
            count += 1
            # print(f"count: {count}")

            # 检查一行中的元素数量是否相同
            if len(mask_values) != len(readout_values) or len(mask_values) != 32:
                raise ValueError("Each line in the mask file and readout file must have 32 bits.")

            modified_readout_values = []
            for i in range(0, len(mask_values)-1, 2):  # 每2位为一组
                # 检查掩码中的两个值是否都为1
                if mask_values[i:i+2] == '11':
                    modified_readout_values.extend([str(special_value)])
                else:
                    # 不需要替换
                    modified_readout_values.extend([readout_values[i], readout_values[i+1]])

            # 将修改后的行写入输出文件
            output_file.write(''.join(modified_readout_values) + '\n')
        print(f"writing rbt file : {output_file_path}")
        
def find_rbd_files(root_path): 
    if not os.path.exists(root_path):
        print(f"Error: The specified path '{root_path}' does not exist.")
        return []
    # rsplit(".rbt", 1) right and only one
    rbd_files = [
        os.path.join(dirpath, filename).rsplit(".rbd", 1)[0]
        for dirpath, _, filenames in os.walk(root_path)
        for filename in filenames if filename.lower().endswith('.rbd')
    ]
    return rbd_files

def main():
    parser = argparse.ArgumentParser(description="Modify rbt trim")

    # Add optional arguments
    parser.add_argument('-m', '--msd_file', type=str, help="Input .msd file path", required=True)
    parser.add_argument('-r', '--rbd_folder', type=str, help="Input .rbd folder path", required=True)
    parser.add_argument('-s', '--special_value', type=str, help="Substitute special value", required=True)

    # Parse arguments
    args = parser.parse_args()

    # Display help if no arguments are provided
    if not any(vars(args).values()):
        parser.print_help()
        return

    print(f"Parameters:")
    print(f"  msd_file: {args.msd_file}")
    print(f"  rbd_folder: {args.rbd_folder}")
    print(f"  special_value: {args.special_value}")

    
    rbd_files = find_rbd_files(args.rbd_folder)

    for path in rbd_files:
        compare_files(args.msd_file, path + ".rbd", path + "_result" + ".rbt", args.special_value)

if __name__ == "__main__":
    main()
